#include "kernel_operator.h"
using namespace AscendC;
#include <type_traits>
#include <typeinfo>
#include <string.h>


constexpr int32_t BUFFER_NUM = 2;      //昇腾双buffer技术


template<typename T>  class KernelScatterElements {
public:
    __aicore__ inline KernelScatterElements() {}


    __aicore__ inline T addFunction(int input_index,int output_index) {
                               
        T tmp = xGm.GetValue(output_index);
        T tmp2 = uGm.GetValue(input_index);
        T tmp_s;
        if constexpr (std::is_same_v<T, uint8_t>) {
            // 类型为 uint8_t 时的代码
            //printf("USE_UINT8\n");
            tmp_s = (T)(static_cast<int>(tmp) + static_cast<int>(tmp2));
        } else {
            // 其他类型时的代码
            //printf("NOT_USE_UINT8\n");
            tmp_s = (T)(static_cast<float>(tmp) + static_cast<float>(tmp2));
            //printf("tmp_s is %f",tmp_s);
        }
        return tmp_s;
    }

    __aicore__ inline T mulFunction(int input_index,int output_index) {
                               
        T tmp = xGm.GetValue(output_index);
        T tmp2 = uGm.GetValue(input_index);
        T tmp_s;
        if constexpr (std::is_same_v<T, uint8_t>) {
            // 类型为 uint8_t 时的代码
            //printf("USE_UINT8\n");
            tmp_s = (T)(static_cast<int>(tmp) * static_cast<int>(tmp2));
        } else {
            // 其他类型时的代码
            //printf("NOT_USE_UINT8\n");
            tmp_s = (T)(static_cast<float>(tmp) * static_cast<float>(tmp2));
        }
        return tmp_s;
    }



    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, int32_t x_size, int32_t indices_size, int32_t updates_size, int32_t x_ndarray[],int32_t indices_ndarray[],int32_t updates_ndarray[], int32_t x_dimensional, int32_t indices_dimensional, int32_t updates_dimensional, int axis, uint8_t reduce)
    {

        this->axis = axis;
        this->reduce = reduce; 
        //printf("axis is %d,reduce is %d\n",this->axis,this->reduce);

        this->x_size =x_size;
        this->x_dimensional = x_dimensional;
        if(this->x_dimensional == 3)
        {
            for(int i = 0; i < 3; i++)
            {
                this->x_ndarray[i] = x_ndarray[i];
            }
        }
        else if(this->x_dimensional == 2)
        {
            this->axis += 1;

            this->x_ndarray[0] = 1;
            for(int i = 1; i < 3; i++)
            {
                this->x_ndarray[i] = x_ndarray[i-1];
            }
        }
        else if(this->x_dimensional == 1)
        {
            this->axis += 2;

            this->x_ndarray[0] = 1;
            this->x_ndarray[1] = 1;
            for(int i = 2; i < 3; i++)
            {
                this->x_ndarray[i] = x_ndarray[i-2];
                //printf(" x_ndarray %d is %f\n",i,this->x_ndarray[i]);
            }
        }
        //printf("x_size is %d, x_dimensional is %dn",this->x_size, this->x_dimensional);

        this->indices_size =indices_size;
        this->indices_dimensional = indices_dimensional;
        if(this->indices_dimensional == 3)
        {
            for(int i = 0; i < 3; i++)
            {
                this->indices_ndarray[i] = indices_ndarray[i];
            }
        }
        else if(this->indices_dimensional == 2)
        {
            this->indices_ndarray[0] = 1;
            for(int i = 1; i < 3; i++)
            {
                this->indices_ndarray[i] = indices_ndarray[i-1];
            }
        }
        else if(this->indices_dimensional == 1)
        {
            this->axis = 2;
            this->indices_ndarray[0] = 1;
            this->indices_ndarray[1] = 1;
            for(int i = 2; i < 3; i++)
            {
                this->indices_ndarray[i] = indices_ndarray[i-2];
                //printf(" indices_ndarray %d is %f\n",i,this->indices_ndarray[i]);
            }
        }
        //printf("indices_size is %d, indices_dimensional is %d\n",indices_size, indices_dimensional);

        this->updates_size =updates_size;
        this->updates_dimensional = updates_dimensional;
        if(this->indices_dimensional == 3)
        {
            for(int i = 0; i < 3; i++)
            {
                this->updates_ndarray[i] = updates_ndarray[i];
            }
        }
        else if(this->indices_dimensional == 2)
        {
            this->updates_ndarray[0] = 1;
            for(int i = 1; i < 3; i++)
            {
                this->updates_ndarray[i] = updates_ndarray[i-1];
            }
        }
        else if(this->indices_dimensional == 1)
        {
            this->updates_ndarray[0] = 1;
            this->updates_ndarray[1] = 1;
            for(int i = 2; i < 3; i++)
            {
                this->updates_ndarray[i] = updates_ndarray[i-2];
                //printf(" updates_ndarray %d is %f\n",i,this->updates_ndarray[i]);
            }
        }
        //printf("updates_size is %d, updates_dimensional is %d\n",updates_size, updates_dimensional);

        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_VAR *>(x), x_size);
        iGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_INDICES *>(indices), indices_size);
        uGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_UPDATES *>(updates), updates_size);
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_VAR *>(y), x_size);

        // ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        // this->coreDataNum = CoreDataNum;
        // this->tileNum = finalTileNum;
        // this->tileDataNum = tileDataNum;
        // this->tailDataNum = TailDataNum;

        // xGm1.SetGlobalBuffer((__gm__ DTYPE_X1*)x1, this->coreDataNum);
        // xGm2.SetGlobalBuffer((__gm__ DTYPE_X2*)x2, this->coreDataNum);
        // yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, this->coreDataNum);

        // pipe.InitBuffer(inQueueX1, BUFFER_NUM, (this->tileDataNum+ 256)  * sizeof(DTYPE_X1));  
        // pipe.InitBuffer(inQueueX2, BUFFER_NUM, (this->tileDataNum+ 256)  * sizeof(DTYPE_X2));  
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, (this->tileDataNum+ 256)  * sizeof(DTYPE_Y));     
        // if constexpr (std::is_same_v<DTYPE_X1, int32_t>)
        // {
        //     pipe.InitBuffer(QueueFloat1, (this->tileDataNum + 256)  * sizeof(float));
        //     pipe.InitBuffer(QueueFloat2, (this->tileDataNum + 256)  * sizeof(float));   
        // }
        // else if constexpr (std::is_same_v<DTYPE_X1, int8_t>)
        // {
        //     pipe.InitBuffer(QueueHalf1, (this->tileDataNum + 256)  * sizeof(half));
        //     pipe.InitBuffer(QueueHalf2, (this->tileDataNum + 256)  * sizeof(half));
        // }
        // pipe.InitBuffer(QueueMask, (this->tileDataNum + 256 )  * sizeof(DTYPE_X1));   //
    }
    __aicore__ inline void Process()
    {
        // int32_t loopCount = this->tileNum;   //手动考虑了双buff
        // this->processDataNum = this->tileDataNum;
        // for (int32_t i = 0; i < loopCount; i++) {
        //     if (i == this->tileNum - 1) {
        //         this->processDataNum = this->tailDataNum;
        //     }
        //     CopyIn(i);
        //     Compute(i);
        //     CopyOut(i);
        // } 
           //1、初始化为0
        
        //DTYPE_UPDATES m_update;

        //memcpy(yGm,0,x_size * sizeof(DTYPE_Y));

        //2、初始化y_ndarray

        int output_index = 0;
        int input_index = 0;

        int m_type = 0;  //1 half 2 uint8_t

        for (int i =0; i <this->x_ndarray[0]; ++i)
        {
            for (int j = 0; j< this->x_ndarray[1]; ++j)
            {
                for(int k = 0; k < this->x_ndarray[2];++k)
                {
                    output_index = i * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
                    //printf("output_index is %d",output_index);
                    yGm.SetValue(output_index, xGm.GetValue(output_index));  
                }                    
            }
        }

        // add
        if(reduce == 'a')   //0 == strcmp(reduce,"add")
        {
            printf("add mode\n");

             for (int i =0; i <this->x_ndarray[0]; ++i)
            {
                for (int j = 0; j< this->x_ndarray[1]; ++j)
                {
                    for(int k = 0; k < this->x_ndarray[2];++k)
                    {
                        if( i < indices_ndarray[0] && j < indices_ndarray[1] && k < indices_ndarray[2])
                        {
                            output_index = i * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
                            input_index = i * indices_ndarray[1] * indices_ndarray[2] + j * indices_ndarray[2] + k;
                            int t = iGm.GetValue(input_index);
                            //printf("t is %d\n",t);
                            if(this->axis == 0)
                            {
                                output_index = t * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
                            }
                            else if (this->axis == 1)
                            {
                                output_index = i* x_ndarray[1] * x_ndarray[2] + t * x_ndarray[2] + k;
                            }
                            else
                            {
                                output_index = i* x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + t;
                            }

                            if(this->updates_size != 1){

                                yGm.SetValue(output_index,addFunction(input_index,output_index)); 
                                //printf("not 1 xGm  is %f, uGm is %f\n",xGm.GetValue(output_index),uGm.GetValue(input_index));
                            }
                            else{ 
                                yGm.SetValue(output_index, uGm.GetValue(0)); 
                                //printf("1 yGm is %f\n",yGm.GetValue(output_index));
                            }
                        }
                    }                   
                }
            }

        }
        // else if (reduce == 'm') //(0 == strcmp(reduce,"multiply"))
        // {
        //     printf("multiply mode\n");

        //     for (int i =0; i <this->x_ndarray[0]; ++i)
        //     {
        //         for (int j = 0; j< this->x_ndarray[1]; ++j)
        //         {
        //             for(int k = 0; k < this->x_ndarray[2];++k)
        //             {
        //                 if( i < indices_ndarray[0] && j < indices_ndarray[1] && k < indices_ndarray[2])
        //                 {
        //                     output_index = i * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
        //                     input_index = i * indices_ndarray[1] * indices_ndarray[2] + j * indices_ndarray[2] + k;
        //                     int t = iGm.GetValue(input_index);
        //                     printf("t is %d\n",t);
        //                     if(this->axis == 0)
        //                     {
        //                         output_index = t * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
        //                     }
        //                     else if (this->axis == 1)
        //                     {
        //                         output_index = i* x_ndarray[1] * x_ndarray[2] + t * x_ndarray[2] + k;
        //                     }
        //                     else
        //                     {
        //                         output_index = i* x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + t;
        //                     }

        //                     if(this->updates_size != 1){
        //                         yGm.SetValue(output_index,mulFunction(input_index,output_index)); 
        //                         printf("not 1 yGm  is %f\n",yGm.GetValue(output_index));
        //                     }
        //                     else{ 
        //                         yGm.SetValue(output_index, uGm.GetValue(0)); 
        //                          printf("1 yGm is %f\n",yGm.GetValue(output_index));
        //                     }
        //                 }
        //             }                   
        //         }
        //     }
        //}
        else
        {
            // self[index[i][j][k]][j][k] = src[i][j][k]  # if dim == 0
            // self[i][index[i][j][k]][k] = src[i][j][k]  # if dim == 1
            // self[i][j][index[i][j][k]] = src[i][j][k]  # if dim == 2
            printf("normal mode\n");
            for (int i =0; i <this->x_ndarray[0]; ++i)
            {
                for (int j = 0; j< this->x_ndarray[1]; ++j)
                {
                    for(int k = 0; k < this->x_ndarray[2];++k)
                    {
                        if( i < indices_ndarray[0] && j < indices_ndarray[1] && k < indices_ndarray[2])
                        {
                            output_index = i * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
                            input_index = i * indices_ndarray[1] * indices_ndarray[2] + j * indices_ndarray[2] + k;
                            int t = iGm.GetValue(input_index);
                            //printf("t is %d\n",t);
                            if(this->axis == 0)
                            {
                                output_index = t * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
                            }
                            else if (this->axis == 1)
                            {
                                output_index = i* x_ndarray[1] * x_ndarray[2] + t * x_ndarray[2] + k;
                            }
                            else
                            {
                                output_index = i* x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + t;
                            }
                            if(this->updates_size != 1){
                                yGm.SetValue(output_index, uGm.GetValue(input_index)); 
                                //printf("not 1 uGm is %f, yGm  is %f\n",uGm.GetValue(input_index),yGm.GetValue(output_index));
                            }
                            else{ 
                                yGm.SetValue(output_index, uGm.GetValue(0)); 
                                //printf("1 yGm is %f\n",yGm.GetValue(output_index));
                            }
                        }
                    }                    
                }
            }
        }

        for (int i =0; i <this->x_ndarray[0]; ++i)
        {
            for (int j = 0; j< this->x_ndarray[1]; ++j)
            {
                for(int k = 0; k < this->x_ndarray[2];++k)
                {
                    input_index = i * x_ndarray[1] * x_ndarray[2] + j * x_ndarray[2] + k;
                    //yGm.SetValue(input_index,0);
                    //printf("final yGm %d is %d\n",input_index, yGm.GetValue(input_index));
                }                    
            }
        }
    
    }


private:
    TPipe pipe;
    //create queue for input, in this case depth is equal to buffer num
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX2;
    //create queue for output, in this case depth is equal to buffer num
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<DTYPE_VAR> xGm;
    GlobalTensor<DTYPE_INDICES> iGm;
    GlobalTensor<DTYPE_UPDATES> uGm;
    GlobalTensor<DTYPE_VAR> yGm;


    //考生补充自定义成员变量
    int32_t x_size;
    int32_t indices_size;
    int32_t updates_size;
    int32_t x_ndarray[20];
    int32_t indices_ndarray[20];
    int32_t updates_ndarray[20];
    int32_t x_dimensional;
    int32_t indices_dimensional;
    int32_t updates_dimensional;
    int axis;
    uint8_t reduce; //[20]


};

extern "C" __global__ __aicore__ void scatter_elements(GM_ADDR x, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelScatterElements<DTYPE_VAR> op;  //,DTYPE_INDICES,DTYPE_UPDATES,DTYPE_VAR
    //补充init和process函数调用内容
    op.Init(x, indices, updates, y, tiling_data.x_size, tiling_data.indices_size, tiling_data.updates_size, tiling_data.x_ndarray, tiling_data.indices_ndarray, tiling_data.updates_ndarray, tiling_data.x_dimensional, tiling_data.indices_dimensional, tiling_data.updates_dimensional, tiling_data.axis,  tiling_data.reduce);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
// call of kernel function
void scatter_elements_do(uint32_t blockDim, void *l2ctrl, void *stream, uint8_t *x, uint8_t *indices, uint8_t *updates, uint8_t *y, 
                   uint8_t *workspace, uint8_t *tiling)
{
    scatter_elements<<<blockDim, l2ctrl, stream>>>(x, indices, updates, y, workspace, tiling);
}
#endif
